# baostock demo for getting history k-line data from server

# 获取历史数据进行回测

import pandas as pd
import baostock as bs
import backtrader as bt
from strategies.strategy_001 import StrategyClass

#### 登陆baostock系统（数据来源） ####
lg = bs.login()
#### 获取历史5分钟K线数据 ####
# query_history_k_data()
fields= "time,code,open,high,low,close"
rs = bs.query_history_k_data("sh.600031", fields,
    start_date='2022-01-01', end_date='2022-06-14',
    frequency="5", adjustflag="2")
data_list = []
while (rs.error_code == '0') & rs.next():
    # 获取一条记录，将记录合并在一起
    data_list.append(rs.get_row_data())
result = pd.DataFrame(data_list, columns=rs.fields)
result.index=pd.to_datetime(result.time, format='%Y%m%d%H%M00000')
data1 = result[['open', 'close', 'high', 'low']].apply(pd.to_numeric)

fields= "date,code,open,high,low,close"
rs = bs.query_history_k_data("sh.600031", fields,
    start_date='2022-01-01', end_date='2022-06-17',
    frequency="d", adjustflag="2")
data_list = []
while (rs.error_code == '0') & rs.next():
    # 获取一条记录，将记录合并在一起
    data_list.append(rs.get_row_data())
result = pd.DataFrame(data_list, columns=rs.fields)
result.index=pd.to_datetime(result.date, format='%Y-%m-%d')
data2 = result[['open', 'close', 'high', 'low']].apply(pd.to_numeric)
bs.logout()


cerebro = bt.Cerebro()
# 设置启动资金
cerebro.broker.setcash(50000.0)
# 设置佣金
cerebro.broker.setcommission(commission=0.001)

cerebro.addstrategy(StrategyClass, half=1000)

dataframe = data1 # 5分钟K线数据
data = bt.feeds.PandasData(dataname=dataframe, timeframe=4)
cerebro.adddata(data)

dataframe = data2 # 日K线数据
data = bt.feeds.PandasData(dataname=dataframe)
cerebro.adddata(data)

print('启动资金：%.2f' % cerebro.broker.getvalue())
cerebro.run(runonce=False)
print('期末资金：%.2f' % cerebro.broker.getvalue())
cerebro.plot(volume=False)[0][0].set_size_inches(9.6,7)